import cv2
import os
from tqdm import tqdm
import numpy as np


def create_mask(imgpath):
    image = cv2.imread(imgpath, cv2.IMREAD_GRAYSCALE)
    _, mask = cv2.threshold(image, 1300, 255, cv2.THRESH_BINARY)
    return mask


def check(image, mask):
    height = mask.shape[0]
    weight = mask.shape[1]
    for row in tqdm(range(height)):  # 遍历高

        for col in range(weight):  # 遍历宽
            if image[row, col, 0] > 240 and image[row, col, 1] > 240 and image[row, col, 2] > 240:
                image[row, col] = [255, 255, 255]
                continue
            elif image[row, col, 0] < 30 and image[row, col, 1] > 230 and image[row, col, 2] < 30:
                image[row, col] = [0, 255, 0]
                continue
            elif image[row, col, 0] < 30 and image[row, col, 1] < 30 and image[row, col, 2] > 230:
                image[row, col] = [0, 0, 255]
                continue
            elif image[row, col, 0] > 230 and image[row, col, 1] < 30 and image[row, col, 2] < 30:
                image[row, col] = [255, 0, 0]
                continue
            elif image[row, col, 0] < 30 and image[row, col, 1] < 30 and image[row, col, 2] > 230:
                image[row, col] = [0, 0, 255]
                continue
            elif image[row, col, 0] > 230 and image[row, col, 1] < 30 and image[row, col, 2] > 230:
                image[row, col] = [255, 0, 255]
                continue
            elif image[row, col, 0] > 230 and image[row, col, 1] < 30 and image[row, col, 2] > 230:
                image[row, col] = [255, 0, 255]
                continue
            elif image[row, col, 0] < 30 and image[row, col, 1] > 230 and image[row, col, 2] > 230:
                image[row, col] = [0, 255, 255]
                continue
            elif image[row, col, 0] < 30 and image[row, col, 1] < 30 and image[row, col, 2] < 30:
                image[row, col] = [0, 0, 0]
                continue
            elif col > 3277:
                image[row, col] = [178, 177, 177]
                continue
            elif mask[row, col] > 200:
                image[row, col] = [178, 177, 177]

    return image


def checkBright(image):
    # averageColors = np.average(image, axis=(0, 1))
    # averageColor = averageColors[0] + averageColors[1] + averageColors[2]
    color_thdLow = 50
    color_thdHigh = 200
    red = []
    green = []
    blue = []
    for row in tqdm(range(image.shape[0])):
        for col in range(image.shape[1]):
            r, g, b = image[row, col]
            if r <= color_thdLow and g <= color_thdLow and b <= color_thdLow:
                continue
            if r >= color_thdHigh and g >= color_thdHigh and b >= color_thdHigh:
                continue
            red.append(r)
            green.append(g)
            blue.append(b)
    aveRed = np.average(np.array(red))
    aveGreen = np.average(np.array(green))
    aveBlue = np.average(np.array(blue))
    return aveRed, aveGreen, aveBlue


if __name__ == '__main__':
    rootpath = "1"
    masksavepath = "mask"
    savepath = "result"
    imgfiles = os.listdir(rootpath)
    brightthreshold = 200
    darkthreshold = 140
    for i in range(0, len(imgfiles)):
        path = os.path.join(rootpath, imgfiles[i])
        print(imgfiles[i])
        if os.path.isfile(path):
            if imgfiles[i].endswith("jpg") or imgfiles[i].endswith("JPG") or imgfiles[i].endswith("bmp"):
                image = cv2.imread(path)
                maskpath = os.path.join(masksavepath, "mask_" + imgfiles[i])
                newname = 'xiufu_' + imgfiles[i].split(".")[0]

                mask = create_mask(path)
                cv2.imwrite(maskpath, mask)

                xiufuImage = check(image, mask)
                cv2.imwrite(os.path.join(savepath, newname + ".jpg"), xiufuImage)
